import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.metrics import mean_squared_error
from sklearn.metrics import mean_absolute_error
from sklearn.metrics import r2_score
# from statsmodels.regression.linear_model import OLS
# from statsmodels.tools import add_constant
from statsmodels.stats.stattools import durbin_watson  

from sklearn.linear_model import LinearRegression

def evaluate_model_df(y_true, y_pred):
    """
    评估回归模型的拟合优度指标，并以表格的形式输出结果。这对于交互式编程
    环境而言，可以很方便地查看结果，也可以轻松地将结果插入到自己的论文中
    去。例如保存为 Excel 文件等。

    输入参数
    ----------
    y_true : array-like
        实际值。
    y_pred : array-like
        预测值。

    返回值
    ----------
    pandas.DataFrame
        各种拟合优度指标的表格。
        
    使用示例
    ----------
    
    简单线性回归模型
    
    >>> X = add_constant(X)
    >>> model = OLS(y, X).fit()
    >>> y_pred = model.predict(X)

    评估模型
    
    >>> evaluate_model(y, y_pred)
    """
    # 计算各项指标
    r2 = r2_score(y_true, y_pred)
    mse = mean_squared_error(y_true, y_pred)
    rmse = np.sqrt(mse)
    mae = mean_absolute_error(y_true, y_pred)
    mape = np.mean(np.abs((y_true - y_pred) / y_true)) * 100

    # 计算调整后的 R2
    n = len(y_true)
    k = 1  # 假设模型只有一个解释变量
    adj_r2 = 1 - (1 - r2) * (n - 1) / (n - k - 1)
    
    # Durbin-Watson 统计量
    residuals = y_true - y_pred
    dw = durbin_watson(residuals)

    # 创建一个包含所有指标的字典
    metrics = {
        'R2': [r2],
        'Adjusted R2': [adj_r2],
        'MSE': [mse],
        'RMSE': [rmse],
        'MAE': [mae],
        'MAPE': [mape],
        'Durbin-Watson': [dw]
    }

    # 转换为 DataFrame
    df_metrics = pd.DataFrame(metrics)

    return df_metrics.T

from sklearn.linear_model import LinearRegression

def prediction_diagnostics(
    y_true, 
    y_pred, 
    ax=None, 
    line_kws=None, 
    scatter_kws=None, 
    fig_kws=None
    ):
    """
    生成预测模型的预测诊断图，包括散点图、灰色的 y=x 虚线以及回归线。

    输入参数
    ----------
    y_true : array-like
        真实值
    y_pred : array-like
        预测值
    ax : matplotlib.axes.Axes, optional
        用于绘图的轴对象。如果未提供则自动创建一个新的轴对象。
    line_kws : dict, optional
        用于调整绘制拟合线样式的关键字参数。
    scatter_kws : dict, optional
        用于绘制散点图的关键字参数。
    fig_kws : dict, optional
        用于创建图形的关键字参数。
        
    返回值
    ----------
    ax : matplotlib.axes.Axes
        经过绘制的轴坐标对象
    """
    if fig_kws is None:
        fig_kws = {"figsize": (5, 5)}
    if ax is None:
        fig, ax = plt.subplots(**fig_kws)
    else:
        fig = ax.figure

    if scatter_kws is None:
        scatter_kws = {
            "color": "none",
            "edgecolor": "blue",
            "label": "Predicted vs. True",
        }

    if line_kws is None:
        line_kws = {
            "color": "gray", 
            "label": "y = x", 
            "linestyle": "--"
        }

    # 绘制散点图
    ax.scatter(y_true, y_pred, **scatter_kws)

    # 绘制 y=x 的虚线
    min_val = min(min(y_true), min(y_pred))
    max_val = max(max(y_true), max(y_pred))
    ax.plot(
        [min_val, max_val], 
        [min_val, max_val], 
        **line_kws)

    # 绘制回归线
    reg = LinearRegression().fit(
        y_true.reshape(-1, 1), y_pred)
    y_fit = reg.predict(
        np.array([min_val, max_val]).reshape(-1, 1))
    ax.plot(
        [min_val, max_val], y_fit, 
        color="red", label="Regression Line")

    # 添加标题和标签
    ax.set_title("Prediction Diagnostics")
    ax.set_xlabel("True Values")
    ax.set_ylabel("Predicted Values")
    ax.legend()

    return ax